# coding: UTF-8
import torch
import numpy as np
from train_eval import init_network, testone
from importlib import import_module
import inputPreprocess
import argparse
import tkinter

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True,
                    help='choose a model: TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer')
parser.add_argument('--embedding', default='pre_trained', type=str, help='random or pre_trained')
parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
args = parser.parse_args()
torch.backends.cudnn.benchmark=True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
torch.backends.cudnn.deterministic = True
vocab = None
config = None
model = None


def run(sentences):
    dataset = 'THUCNews'  # 数据集
    # 嵌入层的预训练词向量设置
    # 搜狗新闻:embedding_SougouNews.npz, 腾讯:embedding_Tencent.npz, 随机初始化:random
    embedding = 'our_embedding.npz'
    if args.embedding == 'random':
        embedding = 'random'
    model_name = args.model  # 'TextRCNN'  # TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer
    if model_name == 'FastText':
        from utils_fasttext import build_dataset, build_iterator
        embedding = 'random'
    else:
        from utils import build_test_dataset, build_iterator, build_test_dataset_from_vocab

    global config, vocab, model
    if config is None:
        x = import_module('models.' + model_name)  # 动态引入模型
        config = x.Config(dataset, embedding)  # 通过module中的Config类定义config
        inputPreprocess.preprocess(sentences, config.test_path)
        vocab, test_data = build_test_dataset(config, args.word)
        config.n_vocab = len(vocab)
        model = x.Model(config).to(config.device)
        if model_name != 'Transformer':
            init_network(model)
    else:
        inputPreprocess.preprocess(sentences, config.test_path)
        test_data = build_test_dataset_from_vocab(vocab, config, args.word)

    test_iter = build_iterator(test_data, config)
    predict_all = testone(config, model, test_iter)
    return [config.class_list[i] for i in predict_all]

if __name__ == '__main__':
    main_window = tkinter.Tk()
    main_window.geometry('1080x640')
    main_window.title('攻击性文本检测工具')

    status = '模型：'+args.model+'  '+'分词：'+str(args.word)
    label1 = tkinter.Label(main_window, text=status)
    label1.grid(row=0, column=1)

    entry = tkinter.Text(main_window, font=16, width=44)
    entry.grid(row=1, column=0)
    label2 = tkinter.Text(main_window, font=16, width=44, state='disabled')
    label2.grid(row=1, column=2)

    def onButton():
        label2.config(state='normal')
        label2.delete(1.0, 'end')
        label2.insert(1.0, '检测中......')
        label2.update()
        results = run(entry.get(1.0, 'end'))
        label2.config(state='normal')
        label2.delete(1.0, 'end')
        label2.insert(1.0, '\n'.join(results))
        label2.config(state='disabled')
    button = tkinter.Button(main_window, text='开始检测', font=18, command=onButton)
    button.grid(row=3, column=1)

    main_window.mainloop()